package com.ruleengine.worker.runner.chain;

import static com.ruleengine.worker.exception.ExceptionMsg.CHAIN_NOT_ENABLED;
import static com.ruleengine.worker.exception.ExceptionMsg.CHAIN_RUNNER_EXIST;
import static com.ruleengine.worker.exception.ExceptionMsg.CHAIN_RUNNER_NOT_EXIST;

import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedDeque;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

import com.github.benmanes.caffeine.cache.Caffeine;
import com.github.benmanes.caffeine.cache.LoadingCache;
import com.github.benmanes.caffeine.cache.RemovalCause;
import com.ruleengine.destination.DestinationManager;
import com.ruleengine.worker.exception.ChainRunnerException;
import com.ruleengine.worker.model.Chain;

import lombok.extern.slf4j.Slf4j;

/**
 * @author lianghanwei on 2024/3/15
 */
@Slf4j
public class CacheChainRunnerManager<T extends ChainRunner> implements ChainRunnerManager {

    private static final int DEFAULT_MAX_CACHED_RUNNER_SIZE = 1000;
    private static final int DEFAULT_MAX_IDLE_TIME_SECONDS = 600;

    private final ExecutorService executor;
    private final ConcurrentMap<String, ConcurrentLinkedDeque<Chain>> chainIdChainMap;
    private final LoadingCache<Chain, ChainRunner> chainRunnerCache;
    private final DestinationManager destinationManager;

    public CacheChainRunnerManager(DestinationManager destinationManager) {
        // TODO-lhw-2024/03/15: executor utils
        this.executor = Executors.newCachedThreadPool();
        this.chainIdChainMap = new ConcurrentHashMap<>();
        this.chainRunnerCache = buildChainRunnerCache();
        this.destinationManager = destinationManager;
    }

    @Override
    public void add(Chain chain) {
        if (exist(chain.getMetadata().getUid())) {
            throw new ChainRunnerException(CHAIN_RUNNER_EXIST);
        }

        getOrAddIfAbsent(chain);
    }

    @Override
    public ChainRunner get(String uid) {
        Chain chain = chainIdChainMap.get(uid) != null
            ? chainIdChainMap.get(uid).peekLast()
            : null;

        return chain != null ? chainRunnerCache.get(chain) : null;
    }

    @Override
    public ChainRunner getOrAddIfAbsent(Chain chain) {
        if (chain.getSpec().getState() != Chain.State.ENABLED) {
            throw new ChainRunnerException(CHAIN_NOT_ENABLED);
        }

        Chain c = chainIdChainMap.computeIfAbsent(chain.getMetadata().getUid(), k -> {
            ConcurrentLinkedDeque<Chain> deque = new ConcurrentLinkedDeque<>();
            deque.add(chain);

            return deque;
        }).getLast();

        return c != null ? chainRunnerCache.get(c) : null;
    }

    @Override
    public boolean exist(String uid) {
        return chainIdChainMap.containsKey(uid);
    }

    @Override
    public void remove(String uid) {
        if (!exist(uid)) {
            throw new ChainRunnerException(CHAIN_RUNNER_NOT_EXIST);
        }

        chainIdChainMap.computeIfPresent(uid, (k, v) -> {
            while(!v.isEmpty()) {
                Chain c = v.pop();
                chainRunnerCache.invalidate(c);
            }

            return null;
        });
    }

    @Override
    public void update(String uid, Chain newChain) {
        if (!exist(uid)) {
            throw new ChainRunnerException(CHAIN_RUNNER_NOT_EXIST);
        }

        chainIdChainMap.computeIfPresent(uid, (k, v) -> update(newChain, v));
    }

    @Override
    public int size() {
        // clear up firstly to get accurate size, e.g. remove expired entries
        chainRunnerCache.cleanUp();

        // chain runner max cache size will not greater 2^31
        return (int) chainRunnerCache.estimatedSize();
    }

    @Override
    public Set<String> uidSet() {
        return chainIdChainMap.keySet();
    }

    @Override
    public void clear() {
        for (String key : chainIdChainMap.keySet()) {
            try {
                remove(key);
            } catch (ChainRunnerException exception) {
                log.error("Remove chain runner failed, key={}", key, exception);
            }
        }
        chainRunnerCache.cleanUp();
    }

    @Override
    public void close() {
        log.info("Start to close chain runner manager.");

        clear();

        log.info("Close chain runner manager successfully.");
    }

    private LoadingCache<Chain, ChainRunner> buildChainRunnerCache() {
        return Caffeine.newBuilder()
            .maximumSize(DEFAULT_MAX_CACHED_RUNNER_SIZE)
            .recordStats()
            .expireAfterAccess(DEFAULT_MAX_IDLE_TIME_SECONDS, TimeUnit.SECONDS)
            .removalListener((Chain key, ChainRunner runner, RemovalCause cause) -> {
                runner.stop();
                runner.close();

                log.info("Remove chain runner from cache successfully, chainUid={}, chainVersion={}.",
                    runner.getUid(), runner.getVersion());
            })
            .executor(executor)
            .build(chain -> {
                ChainRunner runner = RulechainRunnerFactory.createChainRunner(chain, destinationManager);
                if (chain.getSpec().getState() == Chain.State.ENABLED) {
                    runner.start();
                }

                return runner;
            });
    }

    private ConcurrentLinkedDeque<Chain> update(Chain newChain, ConcurrentLinkedDeque<Chain> queue) {
        // 1. create new chain runner
        // 2. add an init(not started) new chain runner to cache
        ChainRunner newRunner = chainRunnerCache.get(newChain);

        // 3. add new chain to map
        queue.add(newChain);

        Chain oldChain = queue.pop();
        ChainRunner oldRunner = chainRunnerCache.getIfPresent(oldChain);

        // 4. stop and close old chain runner
        CompletableFuture.runAsync(() -> {
            if (oldRunner != null) {
                oldRunner.stop();
                oldRunner.close();
            }
            chainRunnerCache.invalidate(oldChain);

            // 5. start new chain runner
            if (newChain.getSpec().getState() == Chain.State.ENABLED) {
                newRunner.start();
            }
        }, executor);

        return queue;
    }
}
